# -*- coding: utf-8 -*-
"""
@Time    : 2025/3/4 20:45 
@Author  : ZhangShenao 
@File    : llm.py 
@Desc    : LLM大模型模块
"""
import os
from typing import List, Dict

import dotenv
from openai import OpenAI, Stream
from openai.types.chat import ChatCompletion, ChatCompletionChunk


class LLM:
    def __init__(self, model_name: str):
        """
        初始化LLM大模型
        :param model_name: LLM大模型名称
        """

        # 加载环境变量
        dotenv.load_dotenv()

        # 创建通义百炼客户端,兼容OpenAI协议
        self._client = OpenAI(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )
        self._model_name = model_name

    def send_msg(self, messages: List[Dict]) -> ChatCompletion | Stream[ChatCompletionChunk]:
        """
        发送消息
        :param messages: 消息列表
        :return: 消息发送结果
        """

        return self._client.chat.completions.create(
            model=self._model_name,
            messages=messages,
        )
